package com.huan.sliding.windowservice;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.core.io.ClassPathResource;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.text.SimpleDateFormat;
import java.util.*;
import java.util.concurrent.TimeUnit;

/**
 * 滑动窗口服务类
 *
 * @author huan.fu
 * @date 2024/11/14 - 21:03
 */
@Component
public class SlidingWindowService {

    private static final Logger log = LoggerFactory.getLogger(SlidingWindowService.class);

    /**
     * 窗口的大小，1s一个窗口
     */
    private static final long WINDOW_SIZE = 1L;
    /**
     * 窗口内请求阈值，最大为5
     */
    private static final long WINDOW_LIMIT = 5;

    private static final RedisScript<Boolean> SLIDING_WINDOW_SCRIPT = RedisScript.of(new ClassPathResource("script/sliding-window.lua"), Boolean.class);

    @Resource
    private RedisTemplate<String, Object> redisTemplate;

    /**
     * 时间窗口 限流
     *
     * @param key key
     * @return true: 允许操作  false: 不允许操作
     */
    public boolean allow(String key) {
        // 当前时间(去掉毫秒位)
        long currentTimeMillis = System.currentTimeMillis() / 1000;
        // 获取窗口的开始时间（当前时间 - 窗口大小）
        long startWindowTime = currentTimeMillis - WINDOW_SIZE;
        // 移除当前窗口开始之前的数据 < startWindowTime
        redisTemplate.opsForZSet().removeRangeByScore(key, 0, startWindowTime);
        // 总计窗口内的总数
        long count = Optional.ofNullable(redisTemplate.opsForZSet().zCard(key)).orElse(1L);

        String startTimeFormat = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss.SSS").format(new Date(startWindowTime * 1000));
        String currentTimeFormat = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss.SSS").format(new Date(currentTimeMillis * 1000));
        log.info("==> 时间窗口 start: {} 和 end: {} 的限制为: {} 当前存在: {} 个元素", startTimeFormat, currentTimeFormat, WINDOW_LIMIT, count);

        // 判断是否超出了窗口限制
        if (count < WINDOW_LIMIT) {
            // 允许方法，添加到当前窗口内 ( 注意： 此处的value不可使用当前时间戳，否则可能统计不准，因为zset中的数据不允许重复 )
            redisTemplate.opsForZSet().add(key, UUID.randomUUID().toString(), currentTimeMillis);
            log.info("** 添加一个元素");
            // 重新设置key的过期时间，非必须，防止这个key删除不掉
            redisTemplate.expire(key, WINDOW_SIZE + 3000, TimeUnit.MILLISECONDS);
            return true;
        }
        log.info("xx 触发限流");
        return false;
    }

    /**
     * 时间窗口 限流
     *
     * @param key key
     * @return true: 允许操作  false: 不允许操作
     */
    public boolean allowForLua(String key) {
        // 当前时间(去掉毫秒位)
        long currentTimeMillis = System.currentTimeMillis() / 1000;
        // 获取窗口的开始时间（当前时间 - 窗口大小）
        long startWindowTime = currentTimeMillis - WINDOW_SIZE;

        String startTimeFormat = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss.SSS").format(new Date(startWindowTime * 1000));
        String currentTimeFormat = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss.SSS").format(new Date(currentTimeMillis * 1000));
        log.info("==> *** 时间窗口 start: {} 和 end: {} 的限制为: {}", startTimeFormat, currentTimeFormat, WINDOW_LIMIT);

        Boolean result = redisTemplate.execute(SLIDING_WINDOW_SCRIPT, Collections.singletonList(key),
                String.valueOf(currentTimeMillis), String.valueOf(WINDOW_LIMIT), String.valueOf(WINDOW_SIZE), UUID.randomUUID().toString());
        if (result != null && !result) {
            log.info("*** 触发限流");
            return false;
        }
        log.info("*** 添加了一个元素");
        return true;
    }
}
